import numpy as np
import matplotlib.pyplot as plt
from mnist_loader import MNISTthin
# 本代码路径     D:\bias\FairDeepLearning-main\dataloaders\data

#数据集保存路径： D:\bias\FairDeepLearning-main\dataloaders\data2
#data2是dataloaders下自己新建的文件夹



# clr_ratio=(0.5,0.5) pos_ratio=(0.5,0.5) sensitiveattr="bck"
data0 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.5, 0.5],egr=0.5,ogr=0.5)
setting0=data0.img
np.save("../data2/balance.npy",setting0)
print("setting0 has been saved")
print("-------------------------------")

# setting1 contains: setting1_1, setting1_2, setting1_3
#clr_ratio=(0.1,0.1) pos_ratio=(0.5,0.5) sensitiveattr="bck"
data1_1 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.1, 0.1],egr=0.5,ogr=0.5)
setting1_1=data1_1.img
np.save("../data2/setting1_1.npy",setting1_1)
print("setting1_1 has been saved")
print("-------------------------------")


#clr_ratio=(0.01,0.01) pos_ratio=(0.5,0.5) sensitiveattr="bck"
data1_2 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.01, 0.01],egr=0.5,ogr=0.5)
setting1_2=data1_2.img
np.save("../data2/setting1_2.npy",setting1_2)
print("setting1_2 has been saved")
print("-------------------------------")



#clr_ratio=(0.001,0.001) pos_ratio=(0.5,0.5) sensitiveattr="bck"
data1_3 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.001, 0.001],egr=0.5,ogr=0.5)
setting1_3=data1_3.img
np.save("../data2/setting1_3.npy",setting1_3)
print("setting1_3 has been saved")
print("-------------------------------")


# setting2 contains: setting2_1, setting2_2
# clr_ratio=(0.1,0.9) pos_ratio=(0.5,0.5) sensitiveattr="bck"
data2_1 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.1, 0.9],egr=0.5,ogr=0.5)
setting2_1=data2_1.img
np.save("../data2/setting2_1.npy",setting2_1)
print("setting2_1 has been saved")
print("-------------------------------")


#clr_ratio=(0.01,0.99) pos_ratio=(0.5,0.5) sensitiveattr="bck"
data2_2 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.01, 0.99],egr=0.5,ogr=0.5)
setting2_2=data2_2.img
np.save("../data2/setting2_2.npy",setting2_2)
print("setting2_2 has been saved")
print("-------------------------------")

#setting3:
#clr_ratio=(0.5,0.5) pos_ratio=(0.9,0.1) sensitiveattr="bck"
data3_1 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.5, 0.5],egr=0.9,ogr=0.1)
setting3_1=data3_1.img
np.save("../data2/setting3_1.npy",setting3_1)
print("setting3_1 has been saved")
print("-------------------------------")

# setting4:
#clr_ratio=(0.5,0.5) pos_ratio=(0.9,0.1) sensitiveattr="color_gy"
data4_1 = MNISTthin(which_set="train",path=r'D:\bias\FairDeepLearning-main\data',clr_ratio=[0.5, 0.5],egr=0.9,ogr=0.1,sensitiveattr="color_gy")
setting4_1=data4_1.img
np.save("../data2/setting4_1.npy",setting4_1)
print("setting4_1 has been saved")
print("-------------------------------")


# ab=np.load(r'D:\bias\FairDeepLearning-main\dataloaders\data2\setting4_1.npy')
# print(ab.shape)
# img = np.swapaxes(ab[1],0,2)
# plt.imshow(img)
# plt.show()
